{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "152d5bf6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: mps\n",
      "Epoch: 0 Loss: 0.27691650390625\n",
      "Epoch: 10 Loss: 0.05615980923175812\n",
      "Epoch: 20 Loss: 0.046093087643384933\n",
      "Epoch: 30 Loss: 0.044670067727565765\n",
      "Epoch: 40 Loss: 0.04451269283890724\n",
      "Epoch: 50 Loss: 0.04433821886777878\n",
      "Epoch: 60 Loss: 0.04425666481256485\n",
      "Epoch: 70 Loss: 0.04423742741346359\n",
      "Epoch: 80 Loss: 0.04422028362751007\n",
      "Epoch: 90 Loss: 0.044215332716703415\n",
      "预测的下一期红球组合: [ 4  9 14 19 24 28]\n",
      "预测的下一期蓝球号码: 8\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "\n",
    "# 1. 数据加载和预处理（包含红球和蓝球）\n",
    "def load_data():\n",
    "    df = pd.read_csv('./lottery_data.csv')\n",
    "    # 反转数据顺序，使最早的日期在前\n",
    "    df = df.iloc[::-1].reset_index(drop=True)\n",
    "    # 提取红球号码\n",
    "    red_balls = df[['red1','red2','red3','red4','red5','red6']].values\n",
    "    # 提取蓝球号码\n",
    "    blue_balls = df['blue'].values.reshape(-1, 1)  # 转换为二维数组\n",
    "    # 合并数据并归一化\n",
    "    combined_data = np.hstack([red_balls, blue_balls])\n",
    "    scaler = MinMaxScaler()\n",
    "    scaled_data = scaler.fit_transform(combined_data)\n",
    "    return scaled_data, scaler\n",
    "\n",
    "# 2. 构建联合预测模型\n",
    "class LotteryPredictor(nn.Module):\n",
    "    def __init__(self, input_size=6, hidden_size=20, output_size=7, num_layers=2, bidirectional=True):\n",
    "        super(LotteryPredictor, self).__init__()\n",
    "        self.bidirectional = bidirectional\n",
    "        self.num_directions = 2 if bidirectional else 1\n",
    "        \n",
    "        # LSTM层处理历史红球数据\n",
    "        self.lstm = nn.LSTM(\n",
    "            input_size=input_size,\n",
    "            hidden_size=hidden_size,\n",
    "            num_layers=num_layers,\n",
    "            bidirectional=bidirectional,\n",
    "            batch_first=False\n",
    "        )\n",
    "        \n",
    "        # 全连接层输出7个数字（6红1蓝）\n",
    "        self.linear = nn.Linear(hidden_size * self.num_directions, output_size)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        out, _ = self.lstm(x)\n",
    "        predictions = self.linear(out)\n",
    "        return predictions\n",
    "\n",
    "# 3. 准备训练数据\n",
    "def prepare_sequences(data, sequence_length=10):\n",
    "    windows = []\n",
    "    labels = []\n",
    "    for i in range(len(data)-sequence_length):\n",
    "        windows.append(data[i:i+sequence_length, :-1])  # 输入只包含红球历史\n",
    "        labels.append(data[i+sequence_length])          # 标签包含红球+蓝球\n",
    "    return np.array(windows), np.array(labels)\n",
    "\n",
    "# 4. 训练函数\n",
    "def train_model():\n",
    "    # 显式检测并设置设备\n",
    "    device = torch.device(\"mps\" if torch.backends.mps.is_available() else \"cpu\")\n",
    "    print(f\"Using device: {device}\")\n",
    "\n",
    "    # 加载数据\n",
    "    scaled_data, scaler = load_data()\n",
    "    X, y = prepare_sequences(scaled_data, sequence_length=10)\n",
    "\n",
    "    # 添加维度调整并确保张量顺序为 (seq_len, batch_size, input_size)\n",
    "    X_tensor = torch.FloatTensor(X).permute(1, 0, 2).to(device)  # (seq_len, batch_size, 6)\n",
    "    y_tensor = torch.FloatTensor(y).to(device)                    # (batch_size, 7)\n",
    "\n",
    "    # 初始化模型\n",
    "    model = LotteryPredictor(\n",
    "        hidden_size=20,\n",
    "        num_layers=2,\n",
    "        bidirectional=True,\n",
    "        output_size=7\n",
    "    ).to(device)\n",
    "    \n",
    "    # 损失函数和优化器\n",
    "    criterion = nn.MSELoss()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "\n",
    "    # 训练模型\n",
    "    epochs = 100\n",
    "    for epoch in range(epochs):\n",
    "        predictions = model(X_tensor)  # (seq_len, batch_size, 7)\n",
    "        loss = criterion(predictions[-1], y_tensor)  # 只计算最后一个时间步的损失\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if epoch % 10 == 0:\n",
    "            print(f'Epoch: {epoch} Loss: {loss.item()}')\n",
    "\n",
    "    return model, scaler, scaled_data\n",
    "\n",
    "# 5. 预测并生成最可能的组合\n",
    "def predict_next_combination(model, scaler, data):\n",
    "    last_sequence = data[-10:]  # 使用最近10期数据预测\n",
    "    \n",
    "    # 提取红球历史（排除最后一列蓝球）\n",
    "    last_red_balls = last_sequence[:, :-1]\n",
    "    \n",
    "    device = torch.device(\"mps\" if torch.backends.mps.is_available() else \"cpu\")\n",
    "\n",
    "    with torch.no_grad():\n",
    "        input_seq = torch.FloatTensor(last_red_balls).unsqueeze(0).permute(1, 0, 2).to(device)\n",
    "        prediction = model(input_seq).squeeze().cpu().numpy()\n",
    "\n",
    "    # 确保 prediction 是 (7,) 形状\n",
    "    if prediction.ndim == 2:\n",
    "        prediction = prediction[-1]\n",
    "\n",
    "    # 反归一化\n",
    "    predicted_values = scaler.inverse_transform(prediction.reshape(1, -1)).astype(int)\n",
    "    \n",
    "    # 拆分红球和蓝球\n",
    "    red_balls = np.clip(predicted_values[0][:6], 1, 33)  # 限制在合法范围\n",
    "    blue_ball = np.clip(predicted_values[0][6], 1, 16)\n",
    "    \n",
    "    return red_balls, blue_ball\n",
    "\n",
    "# 执行训练和预测\n",
    "model, scaler, data = train_model()\n",
    "red_combination, blue_combination = predict_next_combination(model, scaler, data)\n",
    "print(f\"预测的下一期红球组合: {np.sort(red_combination)}\")\n",
    "print(f\"预测的下一期蓝球号码: {blue_combination}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "xquant",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
